InfiniGen: Efficient Generative Inference of Large Language Models with Dynamic KV Cache Management

作者信息

首尔国立大学Jaewoong Sim实验室

链接：[2406.19707] InfiniGen: Efficient Generative Inference of Large Language Models with Dynamic KV Cache Management

摘要

Transformer-based large language models (LLMs) demonstrate impressive performance across various natural language processing tasks. Serving LLM inference for generating long contents, however, poses a challenge due to the enormous memory footprint of the transient state, known as the key-value (KV) cache, which scales with the sequence length and batch size. 【瞬时内存压力特别大】In this paper, we present InfiniGen, a novel KV cache management framework tailored for long-text generation, which synergistically works with modern offloading-based inference systems.【一个可以适用于当前的offload系统的KV Cache管理框架】 InfiniGen leverages the key insight that a few important tokens that are essential for computing the subsequent attention layer in the Transformer can be speculated by performing a minimal rehearsal with the inputs of the current layer and part of the query weight and key cache of the subsequent layer.【insight：哪些是重要的tokens可以被计算出来】 This allows us to prefetch only the essential KV cache entries (without fetching them all), thereby mitigating the fetch overhead from the host memory in offloading-based LLM serving systems.【通过读取重要信息来减少存取】 Our evaluation on several representative LLMs shows that InfiniGen improves the overall performance of a modern offloading-based system by up to 3.00x compared to prior KV cache management methods while offering substantially better model accuracy.

Introduction

Insight

  • 通过上一层的attention计算结果，提前知道在该层时需要预取的kv cache索引
  • 将完整KV Cache offload到CPU上，设计高性能的通信方案
    • 只提取重要的KV Cache信息
    • 提取QK权重局部关键信息

Motivation

  1. 带有offload功能的LLM推理系统

image-20250210110501210

  • a：将KV Cache放置在GPU上速度较快，但会收到内存的限制
  • b：将KV Cache放置在CPU上，因为需要load KV Cache，执行时间非常慢
  • c：即使提前预取KV Cache到GPU，也会因为通信时间大于计算时间而浪费了计算资源
  • d：所以需要一种减少通信的数据量的方案，缩减需要提取的KV Cache数据，实现更好的计算和通信的overlap
  1. Attention计算的局部性和动态性

Attention计算中，只有一部分重要的token是起到作用，它们的Attention分数之和可以非常接近1。H2O就是一个利用Attention计算的局部性加速推理的工作。

image-20241211202039243

方法：

  • H2O方法： H2O是一种KV缓存管理方法，它在每次迭代中通过权重的重要性评估，删除那些认为不重要的KV条目，从而减少KV缓存大小。
  • Optimal方法： Optimal代表一种理想化的情况，假设在每次迭代中，虽然仅保留部分KV缓存条目，但仍然能够以全局视角选择重要的200个token，而不会完全丢弃之前的KV条目。

问题在于，H2O方法假设注意力模式在每次迭代中是静态的，而实际上它是动态变化的。这种假设可能导致重要的KV条目被删除，从而降低模型的生成质量。这就揭示了一个动态性：Iteration级别的动态性

同时，图4同样展示了另外一个动态性：Layer级别的动态性。在Layer0的时候，注意力是比较泛的，所以H2O和基准模型差距比较大；当Layer越来越后的时候，注意力是会慢慢集中在某些token中的，所以H2O和基准模型的差距没那么大。所以Layer0的时候其实是有很大必要探索更多的KV Cache缓存数量的。

image-20241211203924817

图5也同样介绍了Layer层级的动态性问题。就是Layer0时，大部分token关注了大部分K，其注意力分数才能累积到0.9；而在Layer18时，大部分Q其实关注很少的K注意力分数就能累积0.9。

所以，这两种动态性带来一个设计理念：预取数量的动态性，即使在18层，也是有很多需要较多的Key Tokens才能表达的请求。所以对于不同的Query，也应该考虑不同的Key Budget。不能简单地使用Top k算法来提取最重要的k个token的kv cache。

  1. 提前计算出重要的token的索引

image-20250210111944701

image-20241211220857127

Transformer架构中，下一个transformer block的结果是上一个transformer block加上Attn和MLP向量构成的。attn和mlp都经过了layer norm的计算过程，所以离群维度上看变化并不明显。

image-20250210112334659

  • Transformer块的输入向量之间具有高度相似性。

  • InfiniGen利用这一特性，通过上一层的输入（注意力输入）推测当前层的Attention计算结果。

  • SVD在Transformer架构中的使用

Query权重列的差异性

image-20241211222328734

第18层Query权重矩阵的列值分布，其中x轴表示查询矩阵的列索引（即矩阵的列号）。深色的数据为离群值，这幅图表明query权重矩阵中的某些列的作用是明显大于别的列的。

  • 通过乘以矩阵 $A$，重新调整矩阵的方向（“偏斜”），使得查询和键矩阵的注意力权重分布更集中，可以让q1和q2的差距更加明显，然后我们就可以只使用q1来获取近似的attention计算结果。这个矩阵就是V。
  • 使用奇异值分解（SVD）对Transformer模型的查询权重（Query Weight, $W_Q$）和键权重（Key Weight, $W_K$）进行分解。

  • 调整后的矩阵既能减少计算冗余，又不会改变模型的功能和准确性。

    image-20241213144117603

具体公式如下

image-20241227193455369

创新点或贡献

  1. 通过i-1层的attention计算预测第i层attention计算所需要的数据
  2. 在CPU内存阶段就确定了哪些是关键数据，减少了KV Cache数据的传输

具体设计

image-20241211215632743

黄色的组件是Online组件，在推理的同时协同预测和预取关键的KV Cache。蓝色的组件是Offline组件，为了方便预取针对模型权重进行了skewing。

image-20241211225356721

运行流程

image-20241211225647979

image-20250214134335987

image-20250210114313068

推理阶段：动态KV缓存管理

推理阶段分为两个主要子阶段：预填充阶段（Prefill Stage）\和*解码阶段（Decoding Stage）*。

(1) 预填充阶段

预填充阶段负责初始化推理过程，并为后续解码阶段生成关键的部分权重索引：

  1. 输入处理
    • 将输入序列通过模型的初始层（例如Layer 0），生成查询token的初始注意力输入。
  2. 部分权重索引生成（Partial Weight Index Generation, PWIGen）
    • 根据偏斜后的查询权重和键缓存，计算所有列的绝对和（即权重的重要性），并选择最重要的列索引。
    • 选出的部分权重和缓存条目（例如30%的条目）将在解码阶段用于注意力模式的推测。
(2) 解码阶段

解码阶段负责动态推测注意力模式并仅预取必要的KV缓存条目：

  1. 注意力模式推测

    • 在第 $i-1$ 层解码时，使用部分权重和缓存对第 $i$

      层的注意力模式进行推测：

      • 利用第$i-1$层的注意力输入和部分权重计算部分注意力得分（speculated attention scores）。
      • 根据得分的阈值（例如：最大值减去一个偏移量 $\alpha$），选出注意力权重较高的键token。

  2. 关键KV条目预取

    • 仅从CPU缓存池中预取选出的关键KV条目到GPU。
    • 动态调整每层的KV缓存加载量，高效利用PCIe带宽。

KV缓存池管理

KV缓存池管理负责在CPU内存中维护所有生成的KV条目，并在内存压力下进行动态调整：

  1. 缓存存储

    • KV缓存条目默认存储在CPU内存中，GPU仅加载必要的部分。

  2. 内存限制与替换策略

    • 如果CPU内存接近用户定义的限制，系统会移除不常用的KV条目。

    • 替换策略采用轻量化的计数器策略（Counter-based Policy），而不是LRU和FCFS（前者不太好，后者需要一个双链列表，开销比较大）

      • 为每个KV条目维护一个计数器，记录其被访问的频率。
  3. 移除计数器最小的条目以腾出空间，同时更新GPU的缓存条目。

实验评估

先做了准确率的评估

bench mark有量化和H2O

  1. 随着GPU上相关内存的保留得越来越少，InfiniGen准确率都更好

image-20241213113210658

  1. 随着文本长度增长，InfiniGen这种动态的方法可以更贴近Baseline

image-20241213113942674

  1. 文章中提到的Skewing方法（引入A的倾斜率），可以看到引入这种方案可以让模型利用少量的Key和Query列，使得模型结果更贴近Full Cache

image-20241213114427846

  1. KV Cache算法的有效性

  2. 推理延迟方面，InfiniGen从E2E，Batch Szie，Sequence Length和Model Size四个角度完成了实验，体现了InfiniGen关注到的KV数据是非线性增长的。

    image-20241213141742887

    image-20241213141636225

    image-20241213142213227

  3. Sensitivity Study

    1. 分析了文章中Prefill提出的优化中Alpha数字变化的影响
    2. 分析了文章中提出另外一个优化Partial Weight Ratio数字变化的影响

    image-20241213142914058

  4. Overhead

    1. 数据预取
    2. 内存消耗
  5. Long Context Window

背景

image-20241211172138839

LLM KV Cache压力很大。

image-20241211172443917

即使是提前预取，KV Cache传输的数据依旧是非常大的。虽然可以用量化的方法来减轻影响，但线性增长的问题依旧没解决，长文本仍然会遇到这个问题。

为了解决这个问题，过去的工作选择性读取重要KV，来减少KV Cache的读取。

[9] Beidi Chen, Tri Dao, Eric Winsor, Zhao Song, Atri Rudra, and Christopher Ré. Scatterbrain: Unifying sparse and low-rank attention. In Advances in Neural Information Processing Systems (NeurIPS), 2021.

[10] Beidi Chen, Zichang Liu, Binghui Peng, Zhaozhuo Xu, Jonathan Lingjie Li, Tri Dao, Zhao Song, Anshumali Shrivastava, and Christopher Re. Mongoose: A learnable lsh framework for efficient neural network training. In Proceedings of the International Conference on Learning Representations (ICLR), 2021.

[14] Krzysztof Marcin Choromanski, Valerii Likhosherstov, David Dohan, Xingyou Song, Andreea Gane, Tamas Sarlos, Peter Hawkins, Jared Quincy Davis, Afroz Mohiuddin, Lukasz Kaiser, David Benjamin Belanger, Lucy J Colwell, and Adrian Weller. Rethinking attention with performers. In Proceedings of the International Conference on Learning Representations (ICLR), 2021.

[33] Nikita Kitaev, Lukasz Kaiser, and Anselm Levskaya. Reformer: The efficient transformer. In Proceedings of the International Conference on Learning Representations (ICLR), 2020.

[63] Sinong Wang, Belinda Z. Li, Madian Khabsa, Han Fang, and Hao Ma. Linformer: Self-attention with linear complexity. arXiv preprint arXiv:2006.04768, 2020.

还有工作设立了KV Cache的budget，然后使用驱除机制。

[37] Zichang Liu, Aditya Desai, Fangshuo Liao, Weitao Wang, Victor Xie, Zhaozhuo Xu, Anastasios Kyrillidis, and Anshumali Shrivastava. Scissorhands: Exploiting the persistence of importance hypothesis for llm kv cache compression at test time. In Advances in Neural Information Processing Systems (NeurIPS), 2023.

但这些工作都假设了过去某个token，在某一token生成attention时显得不重要了，在别的token生成attention时也不重要。

补充背景

Outliers in Large Language Models

LLM可能会有一些异常值

Singular Value Decomposition

奇异值分解

奇异值分解（Singular Value Decomposition, SVD） 是线性代数中的一种重要矩阵分解方法。它能够将一个矩阵分解为三个特殊矩阵的乘积，是数据分析、机器学习和信号处理等领域的核心工具。

奇异值分解的定义

对于任意一个 $m \times n$的实矩阵 $A$，可以表示为：

$A = U \Sigma V^T $

其中：

  1. $U$ 是一个 $m \times m $ 的正交矩阵，称为左奇异向量矩阵。
  2. $\Sigma $ 是一个 $m \times n$ 的对角矩阵，称为奇异值矩阵，其对角线上的元素为非负数，称为奇异值，且按降序排列。
  3. $V$ 是一个 $n \times n$ 的正交矩阵，称为右奇异向量矩阵。

矩阵 $\Sigma $ 的大小取决于矩阵$A $ 的秩，非零奇异值的个数等于矩阵的秩。

几何意义

  1. 正交变换

    • 矩阵 $A$

      可以被看作是对一个向量进行缩放、旋转的变换。SVD将这个变换分解为三个步骤：

      • $V^T $：对输入进行旋转。
      • $\Sigma $：对旋转后的坐标进行缩放（奇异值对应缩放比例）。
      • $U$ ：对结果进行最终的旋转。

  2. 低秩近似

    • 奇异值分解的奇异值越小，对矩阵贡献越少，因此可以用最大的 $k$ 个奇异值和对应向量近似原矩阵。这种近似广泛用于降维（如PCA）和数据压缩。

SVD的计算方式

给定矩阵 AAA，奇异值分解通过以下方式计算：

  1. 计算 $A^TA$ 和 $AA^T$ 的特征值与特征向量。
  2. $V$ 的列向量是 $A^TA$ 的特征向量。
  3. $U$ 的列向量是 $AA^T$ 的特征向量。
  4. 奇异值是 $A^TA$ 和 $AA^T$ 的特征值的平方根。

应用场景

  1. 降维：
    • 使用奇异值分解提取数据的主要成分（如PCA），将高维数据映射到低维空间。
  2. 矩阵近似：
    • 通过保留最大的奇异值，可以构造低秩矩阵近似，常用于推荐系统、图像压缩等。
  3. 解决最小二乘问题：
    • 通过SVD，可以求解具有约束或欠定问题的最小二乘解。
  4. 特征提取：
    • 在自然语言处理（如LSA）中，用SVD提取文本或词语的潜在语义关系。

直观例子

假设一个 $3 \times 2 $ 矩阵$A$：

$A = \begin{bmatrix} 1 & 0 \ 0 & 2 \ 0 & 0 \end{bmatrix}$

SVD的分解可能是：

$U = \begin{bmatrix} 1 & 0 & 0 \ 0 & 1 & 0 \ 0 & 0 & 1 \end{bmatrix}, \quad \Sigma = \begin{bmatrix} 2 & 0 \ 0 & 1 \ 0 & 0 \end{bmatrix}, \quad V^T = \begin{bmatrix} 0 & 1 \ 1 & 0 \end{bmatrix}$

这个分解将矩阵 AAA 的特性（旋转、缩放）显式表示出来。

图5详解：

x轴（横轴）：每个查询token所需的键token数量

  1. 定义：
    • 表示查询token在计算注意力时，为达到累计注意力权重为0.9所需参与计算的键token的数量。
    • 换句话说，它衡量了每个查询token的注意力模式中，需要考虑多少键token才能覆盖其大部分（90%）的注意力分布。
  2. 含义：
    • 低值（x轴左侧）
      • 查询token只需要少量键token即可满足其注意力需求。
      • 这些token往往高度关注几个特定的键token（注意力模式非常集中）。
    • 高值（x轴右侧）
      • 查询token需要较多的键token来达到0.9的注意力权重。
      • 这些token可能关注较为广泛的上下文信息（注意力模式分散）。
  3. 具体计算方法：
    • 对每个查询token的注意力权重$\text{attention weights} $（通过 $\text{softmax}(QK^T)$ 计算得到）按大小降序排列。
    • 逐一累加这些注意力权重，直到累计值达到0.9。
    • 累加的键token数量即为该查询token所需的键token数量。
  4. 单位：
    • x轴上的数字表示键token的数量，每个bin宽度为16。

y轴（纵轴）：每个bin中查询token的数量

  1. 定义：
    • 表示对于某个所需键token数量的范围（如16、32等），有多少查询token符合这一范围。
    • y轴的值表示查询token的数量分布。
  2. 含义：
    • 较高的y值（某个bin的柱子更高）
      • 说明很多查询token集中在该键token数量范围内。
      • 表明这一层的注意力需求大多数集中在此范围。
    • 较低的y值（某个bin的柱子较矮）
      • 说明在该键token数量范围内的查询token较少，注意力需求较为分散。

x轴和y轴结合的含义

  • x轴和y轴共同描述了查询token在特定层中所需键token数量的分布情况
    • x轴表示查询token需要的键token数量（注意力计算范围）。
    • y轴表示对应范围内的查询token数量。
    • 整体柱状分布反映了该层的注意力模式。

(a) Layer 0 的分布

  • x轴：需要的键token数量较分散，分布较宽（从左到右都有分布）。
  • y轴：每个bin中的查询token数量较为均匀。
  • 含义
    • Layer 0（模型的低层）具有较宽的注意力模式，大量的查询token需要较多的键token来捕获全局上下文。
    • 注意力的权重在许多键token之间分布得较为均匀，因此需要更多键token参与计算。

(b) Layer 18 的分布

  • x轴：需要的键token数量集中在低值区域（分布更窄）。
  • y轴：大多数查询token集中在较少的bin中（较窄范围内的y值较高）。
  • 含义
    • Layer 18（模型的高层）具有更窄的注意力模式，大部分查询token只需要少量键token即可满足注意力需求。
    • 注意力的权重高度集中在少数键token上，说明高层更多地关注局部上下文或特定token。

我如何做这个问题

这个洞见可以引申出其他其他方法吗

该洞见是否可以迁移到其他领域中

KV Cache是瓶颈，从系统层面来减少KV Cache碎片化或者重复的工作很多，典型就是vLLM。这个工作是比较少有的关注减少KV Cache的工作。很有意思。

这个方案可以考虑一下迁移到Star Attention的可行性。

该工作有什么可能可以改进的地方

Q&A

results matching ""

    No results matching ""